package com.paw.distribute.session.resolver;

import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import cn.hutool.core.util.StrUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.session.web.http.CookieSerializer.CookieValue;
import org.springframework.session.web.http.DefaultCookieSerializer;
import org.springframework.session.web.http.HttpSessionIdResolver;

/**
 * @author lvzihai
 * @date 2021/6/18
 **/
@Slf4j
public class TokenCookieHttpSessionIdResolver implements HttpSessionIdResolver {

  private static final String WRITTEN_SESSION_ID_ATTR = TokenCookieHttpSessionIdResolver.class.getName()
      .concat(".WRITTEN_SESSION_ID_ATTR");

  public static final String TOKEN = "token";

  private DefaultCookieSerializer cookieSerializer;

  private boolean useBase64Encoding = false;

  public TokenCookieHttpSessionIdResolver (){
    cookieSerializer = new DefaultCookieSerializer();
    cookieSerializer.setJvmRoute(null);
    cookieSerializer.setUseBase64Encoding(useBase64Encoding);
  }


  @Override
  public List<String> resolveSessionIds (HttpServletRequest request) {
    List<String> list = new ArrayList<>();
    String token = request.getHeader(TOKEN);
    String sessionId = (this.useBase64Encoding ? base64Decode(token) : token);
    if(StrUtil.isNotEmpty(sessionId)){
      list.add(sessionId);
    }
    // 从cookie 中获取
    if(StrUtil.isEmpty(sessionId)){
      list = cookieSerializer.readCookieValues(request);
    }

    return list;
  }

  @Override
  public void setSessionId (HttpServletRequest request, HttpServletResponse response, String sessionId) {
    if (sessionId.equals(request.getAttribute(WRITTEN_SESSION_ID_ATTR))) {
      return;
    }
    request.setAttribute(WRITTEN_SESSION_ID_ATTR, sessionId);
    this.cookieSerializer.writeCookieValue(new CookieValue(request, response, sessionId));
  }

  @Override
  public void expireSession (HttpServletRequest request, HttpServletResponse response) {
    this.cookieSerializer.writeCookieValue(new CookieValue(request, response, ""));
  }

  public boolean isUseBase64Encoding () {
    return useBase64Encoding;
  }

  public void setUseBase64Encoding (boolean useBase64Encoding) {
    this.useBase64Encoding = useBase64Encoding;
  }

  private String base64Decode(String base64Value) {
    try {
      byte[] decodedCookieBytes = Base64.getDecoder().decode(base64Value);
      return new String(decodedCookieBytes);
    }
    catch (Exception ex) {
      log.debug("Unable to Base64 decode value: " + base64Value);
      return null;
    }
  }
}
